import streamlit as st
import os
from milvus.milvus_op import MilvusOP  # Adjust the import according to your module structure
from preprocess.preprocess_data import image_saved_dir
from time import time

# Initialize your Milvus operation object
config = {'ns': 1.0, 'ts': 1.0, 'e': 1.0, 'td': 0.5}
milvus_op = MilvusOP(db_name='state_vector_db', collection_name="hybrid3")

# Streamlit app setup
st.title("Hybrid Search App")

# Sidebar setup
st.sidebar.header("检索配置")

# Use text_area instead of text_input for a larger input box
query = st.sidebar.text_area("输入查询:", placeholder="在此输入您的查询...",
                             height=100, max_chars=500).strip()  # Set height to 100 pixels
query_button = st.sidebar.button("查询", type='primary')
# Configuration sliders
config['ns'] = st.sidebar.slider("NS (文件名BM25匹配)", min_value=0.0, max_value=1.0, value=config['ns'], step=0.01)
config['ts'] = st.sidebar.slider("TS (正文BM25匹配)", min_value=0.0, max_value=1.0, value=config['ts'], step=0.01)
config['td'] = st.sidebar.slider("TD (正文语义匹配)", min_value=0.0, max_value=1.0, value=config['td'], step=0.01)

# Checkbox for additional field matching
config['e'] = 1 if st.sidebar.checkbox("额外字段匹配", value=True) else 0

# Dropdown for limit configuration
limit_options = [3, 5, 10, 20, 50]
limit = st.sidebar.selectbox("选择结果数量:", limit_options, index=1)  # Default to 5

if query_button or query:
    # Perform the search
    stime = time()
    search_res = milvus_op.hybrid_search([query], limit=limit, config=config)
    etime = time()

    # Calculate and display search duration
    search_duration = etime - stime
    st.sidebar.write(f"")

    if search_duration < 0.5:
        color = 'green'
    elif search_duration < 1:
        color = 'yellow'
    else:
        color = 'red'
    st.sidebar.markdown(f"检索耗时: <span style='color:{color};'>{search_duration:.3f}</span> 秒",
                        unsafe_allow_html=True)
    # Display results
    for j, i in enumerate(search_res):
        st.write('-' * 45)
        file_name = f"{i['name']}{i['type']}"
        st.markdown(f"Result {j + 1}: <span style='color:red;'>{file_name}</span>", unsafe_allow_html=True)
        st.write(f"ID: {i['id']}")
        st.write(f"Score: {i['score']} ({i['remark']})")
        # Handle image paths
        if len(i['image_path']) > 0:
            image_path = os.path.join(image_saved_dir, i['image_path'])
            st.image(image_path, caption=image_path)
        else:
            if '|-|' not in i['text']:
                st.markdown(i['text'].replace('\n', '<br>'), unsafe_allow_html=True)
            else:
                st.markdown(i['text'])
else:
    st.sidebar.warning("请输入查询以进行检索。")
